#include <fstream>
#include <cassert>
#include "regalloc.h"
#include "codegen.h"

using namespace std;

void Fixup::setPc(Lir *lir,int pc)
{
    lirToPc[lir] = pc;
}

/// @return -1  the lir doesn't exist
int Fixup::getPc(Lir *lir)
{
    map<Lir*,int>::iterator it;
    it = lirToPc.find(lir);
    if(it != lirToPc.end()){
        return it->second;
    }

    return -1;
}

void Fixup::requireFixup(Lir *lir)
{
    fixups.push_back(lir);
}

void Fixup::doFixup(ostream &os)
{
    vector<Lir*>::iterator it;

    for(it = fixups.begin();it != fixups.end();it++){
        Lir *lir = *it;

        int pos = getPc(lir);
        assert(pos != -1 && "the lir to fixup has no pc");
        os.seekp(pos,ios_base::beg);

        switch(lir->type){
        case JMP:
        {
            AddressOperand *addr = dynamic_cast<AddressOperand*>(lir->arg1);
            Lir *dest_lir = addr->label->lir;

            int dest_pc = getPc(dest_lir);
            assert(dest_pc != -1 && "the destination lir has no pc");
            int rel = dest_pc - (pos + 2); // TODO: support 32-bit jmp, then the size should be 4

            assert(rel < 128 && "rel > 128 means 32-bit jmp is required, not supported now");
            // overwrite the stub
            os.put(0xeb);
            os.put(rel);
            break;
        }
        default:
            break;
        }
    }
}

X86Codegen::X86Codegen()
{
    pc = 0;
    codeStream.open("code.bin");
}

void X86Codegen::emit8(unsigned char d)
{
    codeStream.put(d);
    pc += 1;
}

void X86Codegen::emit32(unsigned int d)
{
    codeStream.write((const char*)&d,sizeof(d));
    pc += 4;
}

unsigned char X86Codegen::makeModRM(char mod,char reg,char rm)
{
    assert(((mod & 0x3) == mod) && "mod is incorrect");
    assert(((reg & 0x7) == reg) && "reg is incorrect");
    assert(((rm & 0x7) == rm) && "rm is incorrect");

    return (unsigned char)( (mod << 6) | (reg << 3) | rm);
}

void X86Codegen::genAdd(Lir *lir)
{
    RegOperand *reg,*rm;
    reg = dynamic_cast<RegOperand*>(lir->dst);
    rm = dynamic_cast<RegOperand*>(lir->arg2);
    unsigned char modRM = makeModRM(0x03,*reg,*rm);
    emit8(0x03);                // r32, r/m32
    emit8(modRM);
}

void X86Codegen::genMove(Lir *lir)
{
    RegOperand *reg,*rm;
    reg = dynamic_cast<RegOperand*>(lir->dst);
    rm = dynamic_cast<RegOperand*>(lir->arg1);
    unsigned char modRM = makeModRM(0x03,*reg,*rm);
    emit8(0x8b);
    emit8(modRM);
}


void X86Codegen::genSub(Lir *lir)
{
    RegOperand *reg,*rm;
    reg = dynamic_cast<RegOperand*>(lir->dst);
    rm = dynamic_cast<RegOperand*>(lir->arg2);
    unsigned char modRM = makeModRM(0x03,*reg,*rm);
    emit8(0x2b);                // r32, r/m32
    emit8(modRM);
}

void X86Codegen::genMul(Lir *lir)
{
    RegOperand *reg,*rm;
    reg = dynamic_cast<RegOperand*>(lir->dst);
    rm = dynamic_cast<RegOperand*>(lir->arg2);
    unsigned char modRM = makeModRM(0x03,*reg,*rm);
    emit8(0x0f);
    emit8(0xaf);                // IMUL r32, r/m32
    emit8(modRM);
}

void X86Codegen::genDiv(Lir *lir)
{
    RegOperand *reg,*rm;
    //reg = dynamic_cast<RegOperand*>(lir->dst);
    rm = dynamic_cast<RegOperand*>(lir->arg2);
    unsigned char modRM = makeModRM(0x03,7,*rm);
    emit8(0xf7);
    emit8(modRM);
}

void X86Codegen::genJmp(Lir *lir)
{
    AddressOperand *addr = dynamic_cast<AddressOperand*>(lir->arg1);
    Lir *dest_lir = addr->label->lir;
    int lir_pc = fixup.getPc(dest_lir);

    if(lir_pc == -1){
        // the destinate lir isn't generated yet, put stubs here
        // TODO: to support jmp to 32-bit offset, we should preserve more bytes
        emit8(0x00);
        emit8(0x00);
        fixup.requireFixup(lir);
    } else {
        // the EIP register contains the address of the instruction following the JMP instruction,
        //so DEST = EIP + REL; EIP = PC + (JMP INSTRUCTION SIZE)

        int rel = lir_pc - (pc + 2) ;
        if(-128 < rel && rel < 128){
            emit8(0xeb);
            emit8(rel);
        } else {
            rel = lir_pc - (pc + 1 + 4);
            emit8(0xe9);
            emit32(rel);
        }
    }
}

void X86Codegen::genPush(Lir *lir)
{
    RegOperand *rm;
    unsigned char reg = 6;
    rm = dynamic_cast<RegOperand*>(lir->arg1);
    unsigned char modRM = makeModRM(0x03,reg,*rm);
    emit8(0xff);
    emit8(modRM);
}

void X86Codegen::genPop(Lir *lir)
{
    RegOperand *rm;
    int reg = 0;
    rm = dynamic_cast<RegOperand*>(lir->arg1);
    unsigned char modRM = makeModRM(0x03,reg,*rm);
    emit8(0x8f);
    emit8(modRM);
}

void X86Codegen::genSpillLoad(Lir *lir)
{
    int offset = *dynamic_cast<ConstantOperand<int>*>(lir->arg1);
    RegOperand *reg = dynamic_cast<RegOperand*>(lir->dst);
    unsigned char mod,rm;

    emit8(0x8b);                // MOV r32, r/m32
    if(offset < 128){
        rm = 0x03;              // [EBX] + disp8
        mod = 0x01;
        emit8(makeModRM(mod,*reg,rm));
        emit8(offset);
    } else {
        rm = 0x05;              // [EBX] + disp32
        mod = 0x02;
        emit8(makeModRM(mod,*reg,rm));
        emit32(offset);
    }
}

void X86Codegen::genSpillStore(Lir *lir)
{
    int offset = *dynamic_cast<ConstantOperand<int>*>(lir->dst);
    RegOperand *reg = dynamic_cast<RegOperand*>(lir->arg1);
    unsigned char mod,rm;

    emit8(0x89);                // MOV r/m32,r32
    if(offset < 128){
        rm = 0x03;              // [EBX] + disp8
        mod = 0x01;
        emit8(makeModRM(mod,*reg,rm));
        emit8(offset);
    } else {
        rm = 0x05;              // [EBX] + disp32
        mod = 0x02;
        emit8(makeModRM(mod,*reg,rm));
        emit32(offset);
    }
}

void X86Codegen::genLoadConst(Lir *lir)
{
    // TODO: fake implement here, use the actual value in constant pool
    RegOperand *reg = dynamic_cast<RegOperand*>(lir->dst);
    ConstantOperand<int> *index = dynamic_cast<ConstantOperand<int>*>(lir->arg1);

    emit8(0xb8+*reg);            // MOV r32,imm32
    emit32(*index);
}

void X86Codegen::genOr(Lir *lir)
{
    ConstantOperand<int> *imm = dynamic_cast<ConstantOperand<int>*>(lir->arg1);
    RegOperand *rm = dynamic_cast<RegOperand*>(lir->dst);
    unsigned char reg = 0x01;
    unsigned char modRM = makeModRM(0x03,reg,*rm);

    emit8(0x81);                // OR r/m32,imm32
    emit8(modRM);
    emit32(*imm);
}

void X86Codegen::genCode(CFG *cfg)
{
    RegAlloc *regalloc = cfg->regalloc;

    for(CFG::iterator it = cfg->begin();it != cfg->end();it++){
        BasicBlock *bb = *it;
        LirQueue &q = regalloc->getLirsFor(bb);

        for(LirQueue::iterator lit = q.begin();lit != q.end();lit++){
            Lir *lir = *lit;

            // record the pc for this lir for fixup progress
            fixup.setPc(lir,pc);

            switch(lir->type){
            case ADD:
                genAdd(lir);
                break;
            case SUB:
                genSub(lir);
                break;
            case MUL:
                genMul(lir);
                break;
            case DIV:
                genDiv(lir);
                break;
            case MOV:
                genMove(lir);
                break;
            case JMP:
            case JNE:
                genJmp(lir);
                break;
            case PUSH:
                genPush(lir);
                break;
            case POP:
                genPop(lir);
                break;
            case SPILL_LOAD:
                genSpillLoad(lir);
                break;
            case SPILL_STORE:
                genSpillStore(lir);
                break;
            case LOADCONST:
                genLoadConst(lir);
                break;
            case ORri:
                genOr(lir);
                break;
            }
        }
    }

    fixup.doFixup(codeStream);
}









